import numpy as np
from collections import Counter

import pandas as pd


class RDI(object):
    def __init__(self, N=52, save_dir='outputs'):
        '''
        :param N(int): 养护路段总长度(Km)，向上取整
        '''
        super().__init__()
        # 路面车辙性能预测模型参数表
        self.a = np.array([[0.0458, 0.0470, 0.0375, 0.0845],
                           [0.0417, 0.0452, 0.0411, 0.0619],
                           [0.0631, 0.0631, 0.0506, 0.0613],
                           [0.0726, 0.0810, 0.0685, 0.0726],
                           [0.0851, 0.0893, 0.0768, 0.0643],
                           [0.0887, 0.0929, 0.0804, 0.0637],
                           [0.1030, 0.1030, 0.0905, 0.0738],
                           [0.1095, 0.0195, 0.0970, 0.0804]])
        self.b = np.array([[0.3101, 0.3089, 0.5518, 0.2845],
                           [0.3345, 0.2929, 0.4339, 0.3095],
                           [0.1036, 0.1036, 0.3006, 0.2131],
                           [0.0202, -0.0381, 0.1589, 0.1298],
                           [-0.0815, -0.0940, 0.1030, 0.1905],
                           [-0.1113, -0.1238, 0.0732, 0.1899],
                           [-0.2637, -0.2304, -0.0333, 0.0833],
                           [-0.3262, -0.2929, -0.0958, 0.0208]])
        self.a0 = np.array([1.0])
        self.a1 = np.array([3.0])

        self.N = N
        self.rds, self.ri = self.statistics_rds()  # 统计养护路段
        self.rds_list = []
        self.rds_list.append(self.rds)
        self.save_rds(save_dir)

    def statistics_rds(self):
        # 本年路面使用性能指标(RD，单位mm)
        rds = np.random.randint(2, 15, (self.N, 1)) * 1.0  # 随机生成每个路段的车辙深度(RD)
        # 养护概率 (测试)
        # H = np.ceil(np.abs(np.random.normal(3, 1.8, size=(N, 100))))
        ratio = np.random.normal(0.2, 0.025)  # 正态分布 超90%概率在[46, 54]区间
        i = np.random.choice(self.N - 1, np.floor(self.N * (1 - ratio)).astype(np.int_), replace=False)
        rds[i] = 0  # i 路段不养护
        # rds = np.expand_dims(rds, axis=1).repeat(self.ND, axis=1).squeeze()
        return rds, i

    def update_rds(self, var):
        self.rds = var

    def save_rds(self, dir_):
        df = pd.DataFrame(self.rds.squeeze(), columns=['rd'])
        # print(df)
        df.to_csv('{}/养护指标.csv'.format(dir_))

    def decay_rd(self, t, rd):
        delta = 2 * self.a * t + self.b
        if rd == 0:
            return rd * delta
        return rd - delta

    def RD(self, t, j, b):
        '''
        :param t: 养护年
        :param j: 养护类别
        :param b: 路段交通量
        :return:
        '''

        jb = np.hstack((j, b))
        rds = []
        for rd in self.rds:
            xs = []
            for i, x in enumerate(rd):
                xs.append(self.decay_rd(t, x)[tuple(jb[i])])
            rds.append(xs)

        # 计算衰减车辙性能(RD)
        rds_ = np.array(rds)
        # rds_ = np.array(rds)

        # rds_ = np.stack([self.decay_rd(t, rd)] for rd in self.rds).squeeze()
        # for rd in rds_:
        #     xx = rd[j, b]

        # 更新当前年t路面性能指标(RD)
        # for i, x in enumerate(rds_):
        #     self.rds[i, 0] = x[j, b]
        # self.rds shape:
        #     (N, X) N 路段总数， X 种群总数
        # self.rds = np.stack([rd[j, b]] for rd in rds_).squeeze()
        # self.rds = rds_[:, :, j, b].squeeze()

        # 计算当前年t综合路面性能(RD)
        res = []
        for rd in rds_:
            x1 = np.stack([self.a * t ** 2 + self.b * t + x] for x in rd).squeeze()
            ds = []
            for i, x in enumerate(x1):
                ds.append(x[tuple(jb[i])])
            # print(ds)
            res.append(ds)

        # res = np.stack([self.a * t ** 2 + self.b * t + rd[j, b]] for rd in self.rds).squeeze()
        # self.total_rds = np.stack((self.total_rds, np.array(res)))
        # self.rds_list.append(np.array(res))

        return np.array(res)  # shape(N, X)

    def rd_new(self, t, js, b=0):
        '''
        :param t: 养护年
        :param j: 养护类别
        :param b: 路段交通量
        :return:
        '''

        def convert(rds, jx, b):
            dds = dict(enumerate(jx.flatten(), 0))
            i = np.array(list(zip(dds.keys(), dds.values())))
            return rds[:, :, b].squeeze()[i[:, 0], i[:, 1]]

        rds_ = np.stack([self.decay_rd(t, rd)] for rd in self.rds).squeeze()

        # for jx in js:
        #     print(jx)
        #
        # for jx, bx in zip(js, b):
        #     print(jx)
        #     convert(rds_, js, b)
        #
        # res = np.stack([convert(rds_, jx, 0)] for jx in js).squeeze()
        res = np.stack([convert(rds_, jx, b)] for jx in js).squeeze()
        return res

    def rd_result(self, t, j, b):
        '''
        :param t: 养护年
        :param j: 养护类别
        :param b: 路段交通量
        :return:
        '''

        def convert(rds, jx, bx):
            dds = dict(enumerate(jx.flatten(), 0))
            i = np.array(list(zip(dds.keys(), dds.values())))
            return rds[:, :, bx].squeeze()[i[:, 0], i[:, 1]]

        rds_ = np.stack([self.decay_rd(t, rd)] for rd in self.rds).squeeze()

        res = convert(rds_, j, b)

        # res = np.stack([convert(rds_, jx, bx)] for jx, bx in zip(js, b)).squeeze()
        return res

    def RDI(self, rd):
        rdi = 100.0 - self.a0 * rd
        return rdi

    def calc_maintenance_type(self, t, j, b=0):
        def convert(x):
            x[x <= 2.] = 0
            x[(x > 2.) & (x <= 3.)] = 1
            x[(x > 3.) & (x <= 5.)] = 2
            x[(x > 5.) & (x <= 8.)] = 3
            x[(x > 8.) & (x <= 10.)] = 4
            x[(x > 10.) & (x <= 12.)] = 5
            x[(x > 12.) & (x <= 15.)] = 6
            x[x > 15.] = 7
            return x

        rds = self.rd_new(t, j, b)
        # print(rds)
        res = np.array(convert(rds), dtype=np.int_)
        return res

    def calc_result_maintenance_type(self, t, j, b=0):
        def convert(x):
            x[x <= 2.] = 0
            x[(x > 2.) & (x <= 3.)] = 1
            x[(x > 3.) & (x <= 5.)] = 2
            x[(x > 5.) & (x <= 8.)] = 3
            x[(x > 8.) & (x <= 10.)] = 4
            x[(x > 10.) & (x <= 12.)] = 5
            x[(x > 12.) & (x <= 15.)] = 6
            x[x > 15.] = 7
            return x

        rds = self.rd_result(t, j, b)
        # print(rds)
        res = np.array(convert(rds), dtype=np.int_)
        return res


def func(b, js):
    dds = dict(enumerate(js[0].flatten(), 0))
    dict_tuple = list(zip(dds.keys(), dds.values()))

    for t in range(T):
        rds_ = np.stack([rdi.decay_rd(t, rd)] for rd in rdi.rds).squeeze()


if __name__ == '__main__':
    T = 5
    rdi = RDI()
    # # j = np.random.randint(0, 8, (100, 1))
    # b = np.random.randint(0, 4, (100, 1))
    # vs = np.random.randint(0, 7, (100, 52))
    # # b = vs[:, [0]]
    # js = vs[:, 1:]

    var = np.array(
        [[0, 5, 5, 7, 3, 3, 3, 3, 7, 3, 3, 5, 7, 7, 7, 5, 7, 3, 7, 5, 3, 5, 5, 7, 3, 7, 6, 3, 7, 7, 7, 7, 4, 6, 7, 3, 7,
          5, 7, 7, 7, 7, 3, 5, 6, 5, 7, 6, 7, 7, 5, 5, 3]]).reshape(-1, 1)
    b = var[0]
    js = var[1:]
    outs = []
    for t in range(T):
        # rds_ = np.stack([rdi.decay_rd(t, rd)] for rd in rdi.rds).squeeze()
        # print(rds_)
        out = rdi.calc_result_maintenance_type(t, b, js)
        d = dict(Counter(out))
        outs.append(dict(sorted(d.items(), key=lambda d1: d1[0])))
        print(out)

    rds = np.array(pd.read_csv('养护指标.csv', usecols=['rd']))
    print(rds)
    #
    # # func(b, js)
    # # dds = dict(enumerate(js[0].flatten(), 0))
    # # dict_tuple = zip(dds.keys(), dds.values())
    # # for i in dict_tuple:
    # #     print(i)
    # # for i, x in enumerate(js):
    # #     tx = tuple((x, b[i]))
    # #     print(x)
    # for t in range(T):
    #     # rd = rdi.RD(t, j, b)
    #     rds_ = np.stack([rdi.decay_rd(t, rd)] for rd in rdi.rds).squeeze()
    #
    #     # res = np.stack([convert(rds_, jx, bx)] for jx, bx in zip(js, b)).squeeze()
    #
    #     # for jx, bx in zip(js, b):
    #     #     print(jx)
    #     #     print(bx)
    #     #     dds = dict(enumerate(jx.flatten(), 0))
    #     #     dict_tuple = list(zip(dds.keys(), dds.values()))
    #     #     b = np.array(list(zip(dds.keys(), dds.values())))
    #     #     rds_[:, :, bx].squeeze()[b[:, 0], b[:, 1]]
    #
    #     rd = rdi.decay_rd(t, rdi.rds)
    #     print(rd)
    #     rdi.calc_maintenance_type(t)
